package baconsnake;

/*
 * BaconSnake (c) Arnab Nandi, http://arnab.org/blog/baconsnake
 */
import java.io.IOException;
import java.util.HashMap;

import org.apache.pig.EvalFunc;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.DataType;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.util.WrappedIOException;
import org.python.core.PyFunction;
import org.python.core.PyList;
import org.python.core.PyString;
import org.python.util.PythonInterpreter;

public class EvalFunction extends EvalFunc<String> {

	static HashMap<String, PyFunction> pyFunc = null;

	public static void init(String functionName, String pythonCode) {
		if (pyFunc == null) {
			pyFunc = new HashMap<String, PyFunction>();
		}

		PythonInterpreter interp = new PythonInterpreter();
		interp.exec(pythonCode);
		pyFunc.put(functionName, (PyFunction) interp.get(functionName,
				PyFunction.class));
	}

	public String exec(Tuple input) throws IOException {
		if (input == null || input.size() != 3)
			return null;
		String functionName = null;
		String pythonCode = null;

		try {
			functionName = (String) input.get(0);
			pythonCode = (String) input.get(1);
		} catch (Exception e) {
			e.printStackTrace();
			throw WrappedIOException.wrap(
					"Caught exception parsing BaconSnake call ", e);
		}

		try {

			if (pyFunc == null || !pyFunc.containsKey(functionName)) {
				EvalFunction.init(functionName, pythonCode);
			}
		} catch (Exception e) {
			e.printStackTrace();
			throw WrappedIOException.wrap(
					"Caught exception initializing BaconSnake function ", e);
		}

		try {
			// We accept Strings
			if (input.getType(2) == DataType.CHARARRAY) {
				String str = (String) input.get(2);
				PyString pi = (PyString) pyFunc.get(functionName).__call__(
						new PyString(str));
				return pi.toString();
			}
			// or DataBags
			else if (input.getType(2) == DataType.BAG) {
				PyList pl = toList((DataBag) input.get(2));
				PyString pi = (PyString) pyFunc.get(functionName).__call__(pl);
				return pi.toString();
			} else {
				return null;
			}

		} catch (Exception e) {
			e.printStackTrace();
			throw WrappedIOException.wrap(
					"Caught exception processing input row ", e);
		}
	}

	// Recursively add tuples to a nested list of strings
	PyList toList(Tuple tup) throws ExecException {
		PyList pl = new PyList();
		for (int i = 0; i < tup.size(); i++) {
			if (tup.getType(i) == DataType.CHARARRAY) {
				pl.add(new PyString((String) tup.get(i)));
			} else if (tup.getType(i) == DataType.TUPLE) {
				pl.add(toList((Tuple) tup.get(i)));
			}
		}
		return pl;
	}

	// A DataBag can only hold tuples
	PyList toList(DataBag dbag) throws ExecException {
		PyList pl = new PyList();
		for (Tuple tup : dbag) {
			pl.add(toList(tup));
		}
		return pl;
	}
}
